package com.varian.tool.tree;

import cn.hutool.core.collection.CollUtil;
import com.varian.tool.Func;

import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * @author ben
 * @since 2024/7/4
 */
public class TreeNodeManager<T extends INode<T>> {

    private final Map<Long, T> nodeMap;
    private final Set<Long> parentIds;
    private final List<T> nodeList;

    public TreeNodeManager(List<T> nodeList) {
        this.parentIds = CollUtil.newHashSet();
        this.nodeMap = Func.unmodifiableMap(nodeList, INode::getId);
        this.nodeList = Func.unmodifiableList(nodeList);
    }

    public List<T> build() {
        if (CollUtil.isNotEmpty(nodeList)) {
            resolve();
            return nodeMap.values().stream()
                    .filter(node -> node.getParentId().equals(Func.TREE_ROOT_ID) || parentIds.contains(node.getId()))
                    .sorted(INode::compareTo)
                    .toList();
        }
        return Collections.emptyList();
    }

    private void resolve() {
        nodeList.forEach(node -> {
            if (nodeMap.containsKey(node.getParentId())) {
                addChild(nodeMap.get(node.getParentId()), node);
            } else {
                parentIds.add(node.getId());
            }
        });
    }

    private void addChild(T parent, T node) {
        if (parent.getChildren() == null) {
            parent.setChildren(CollUtil.newArrayList());
        }
        parent.getChildren().add(node);
    }
}
